
import os
from torch.utils.data import Dataset
from skimage import io
from skimage.color import rgb2gray
import torch
import matplotlib.pyplot as plt
from torchvision import transforms
from torch.utils.data import DataLoader

labels = []
for i in range(10):
    labels.append(48 + i)
for i in range(26):
    labels.append(97 + i)
for i in range(26):
    labels.append(65 + i)
print(labels)
# print(ord('A'))

class VerCodeDataset(Dataset):
    def __init__(self, image_dir):
        cwd = os.getcwd()
        self.data = []
        self.label = []

        image_dir = os.path.join(image_dir, 'img_train')
        print(f'image_dir:{image_dir}')

        l = os.listdir(image_dir)
        for d in l:
            if d == '.DS_Store':
                continue
            fs = os.listdir(os.path.join(image_dir, d))
            for f in fs:
                img_path = os.path.join(image_dir, d, f)
                t = torch.from_numpy(io.imread(img_path)).float() / 255

                norl = transforms.Normalize(t.mean(), t.std())
                data_norl = norl(t.reshape(1, 27, 26))
                # print(f'data_norl:{data_norl}')
                self.data.append(data_norl)
                # self.label.append(d)
                # print(d)
                label_index = labels.index(ord(d))
                # print(f'label_index:{label_index},d:{d},ord(d):{ord(d)}')
                self.label.append(label_index)

    def __len__(self):
        ''' 返回数据集的长度 '''
        return len(self.data)

    def __getitem__(self, item):
        ''' 数据集获取时的操作 '''
        # print(f'item:{item}')
        return {"data": self.data[item], "label": self.label[item]}

def trainloader(bs):
    ds = VerCodeDataset(image_dir='./learn/')
    return DataLoader(ds,batch_size=bs)

def testloader():
    ds = VerCodeDataset(image_dir="./image_test/")
    return DataLoader(ds,batch_size=5)

if (__name__ == "__main__"):
    tl = trainloader(5)
    # print(tl)
    for step,i in enumerate(tl):
        print(f'step:{step};i:{i}')
        print(f'i["data"]:{i["data"]}')
        print(f'i["label"]:{i["label"]}')
        exit(0)
